package ribbon.consumer.selfrule;

/**
 * @Author:RenPu
 * @Date: 2020/8/25 16:49
 * @Version: 1.0
 * @description: 使用一致性hash算法，手写一个负载均衡的策略
 */


import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;
import lombok.NoArgsConstructor;
import org.apache.tomcat.util.http.fileupload.RequestContext;
import org.omg.CORBA.ServerRequest;
import org.springframework.util.CollectionUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.io.UnsupportedEncodingException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Collections;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;

/**
 * 一致性哈希：
 * 脑图就是一个封闭的圆环，分部了多个服务器，根据一个request的特征亮
 * （如：参数，url，字符串等等），通过hash算法生成一定的摘要也就是int值，
 * 分布在这个圆环上，然后根据这int的值，按照顺时针去寻找距离这个int值，相近的服务
 * 因为request的特征量，都是固定的。所以访问的服务也是同一个机器
 */

@NoArgsConstructor
public class MyRule extends AbstractLoadBalancerRule implements IRule {
    @Override
    public void initWithNiwsConfig(IClientConfig iClientConfig) {

    }

    @Override
    public Server choose(Object o) {

        //通过RequestContextHolder上下文获取，request对象
        HttpServletRequest request = ((ServletRequestAttributes)
                RequestContextHolder.getRequestAttributes()).getRequest();


        //获取uri为例如：www.baidu.com/ninini,斜杠后面的值
        String uri = request.getServletPath() + "?" + request.getQueryString();


        //根据request的特征亮，获取唯一的hashcode的值，以及根据LoadBalancer方法获取所有服务的节点
        return route(uri.hashCode(), getLoadBalancer().getAllServers());
    }


    //将服务通过一致hash，散落在不同的位置上
    public Server route(int hashId, List<Server> addressList) {


        //预防攻击编程
        if (CollectionUtils.isEmpty(addressList)) {
            return null;
        }

        TreeMap<Long, Server> treeMap = new TreeMap<>();

        addressList.stream().forEach(e -> {

            for (int i = 0; i < 8; i++) {
                Long hash = hash(e.getId() + i);
                treeMap.put(hash, e);
            }
        });

        Long hash = hash(String.valueOf(hashId));
        SortedMap<Long, Server> last = treeMap.tailMap(hash);

        //当requets url的hash值，大于任意一个服务器对应的hashKey,则取treeMap的第一个节点
        if (last.isEmpty()) {
            treeMap.firstEntry().getValue();

        }

        return last.get(last.firstKey());
    }


    //封装使用md5实现一致性hash算法
    public Long hash(String key) {

        MessageDigest md5;

        try {
            md5 = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }

        byte[] bytes = null;
        try {
            bytes = key.getBytes("UTF-8");
        } catch (UnsupportedEncodingException e) {
            e.printStackTrace();
        }

        md5.update(bytes);
        byte[] digest = md5.digest();

        long hashCode = ((long) (digest[2] & 0xFF << 16)) | ((long) (digest[1] & 0xFF << 8)) | ((long) (digest[0] & 0xFF));

        return hashCode & 0xffffffffL;
    }
}
